"""
相较于 draw 模块，更高级的画图模块，整合、抽象程度更高

start at 2024-02-26 by lilidream
"""

from typing import List, Tuple, Union, Callable
from matplotlib import pyplot as plt
from matplotlib.figure import Figure
from matplotlib.gridspec import GridSpec
from matplotlib.ticker import MultipleLocator
from matplotlib.contour import QuadContourSet
from matplotlib.quiver import Quiver
import cartopy.crs as ccrs
from cartopy.mpl.geoaxes import GeoAxes
from cartopy.mpl.gridliner import Gridliner
import yaml
import xarray as xr
from . import const as _C
import numpy as np
from tqdm import tqdm

class Painter:
    """
    高级画图类

    Example:
    ```
    p = Painter(2, 2, mapRange=mapRange, figZoom=2)

    def draw(ax, row, col, i, j):
        print(i)
        d = clusterData.sel(offset=0, type=trackType[i])
        C = p.contourf(ax, d.q*1000, preset='q')
        p.quiver(ax, d.u, d.v, lowRes=10)
        return C
        
    p.for_each_axes(draw)
    p.for_each_label(lambda r,c,i,j : "Fig "+str(i))
    p.colorbar(label="SH")
    p.quiverkey()
    p.fig_title("test")
    ```
    """

    figSize: tuple
    """fig的figSize"""

    fig: Figure
    """matplotlib 的 Figure 实例"""

    gridSpec: GridSpec
    """多图的网格布局"""

    axes: List[List[Union[plt.Axes, GeoAxes]]] = [[None]]
    """全部子图的对象列表(二维)，如果 customAxes 为 False，则为规整的网格子图Axes对象列表"""

    config: dict
    """配置文件的内容"""

    _contourf: QuadContourSet
    """用于保存最后一个画的 contourf 对象"""

    _quiver: Quiver
    """用于保存最后一个画的 quiver 对象"""

    def __init__(self, rows: int=1, cols: int=1, figSize: tuple=None,
                 figureRatio: float=None, mapRange: list=None,
                 customAxes: bool=False, customMap: bool=False, 
                 config:Union[str, dict]="default", **kwargs) -> None:
        """
        创建一个 小画家(Painter) 对象，输入基本的画图参数

        :param rows: 子图的行数。
        :param cols: 子图的列数。
        :param figSize: 画图的大小，格式为 (width, height)，如果不设置则根据其他参数自动计算。
        :param figureRatio: 子图的宽高比，如果是地图则按照地图计算，忽略此参数。
        :param mapRange: 地图的范围，格式为 [lon1, lon2, lat1, lat2]，
            如果设置了此参数，则每一张子图都是地图，除非设置 customAxes 为 True。
        :param customAxes: 是否自定义子图，如果为 False，则生成规整的网格子图，
            如果要绘制大小、位置不同的子图，则设置为 True，并使用 add_axes() 方法添加子图。
        :param customMap: 是否自定义地图，如果为 False，则使用默认的 PlateCarree 投影地图，
            并绘制海岸线和网格线，如果想自定义则使用 add_axes() 方法并传入 mapRange 参数，
            设置为 GeoAxes，然后在用 draw_coastlines() 等方法绘制地图。
        :param config: yaml 配置文件路径，不传入则使用默认配置文件，其他传入参数会覆盖默认配置。
            或者直接传入配置文件的 dict 对象。

        其他参数(配置文件参数，传入任意与配置文件中相同的参数名即可覆盖配置文件中的参数，下面列举部分)：
        :param figZoom: 子图的基本高度。
            如果 figSize 未设置，则每一张子图的高度约为 (figZoom * dpi) px。
        :param dpi: 画图的 dpi。
        :param sub_adj: 整个图的绘图范围，格式为 [left, bottom, right, top]，
            同 matplotlib.figure.Figure.subplots_adjust()
        :param wspace: 子图的宽间隔。
        :param hspace: 子图的高间隔。
        :param widthShrink: 用于收缩整个图的宽度，因为地图不一定按比例，所以手动调整此参数。
        """
        # 加载配置文件
        if isinstance(config, str):
            self._load_config(
                (_C._LIB_PATH + "config/painter_default.yaml") if config == "default" else config
                )
        elif isinstance(config, dict):
            self.config = config
        else:
            raise ValueError("config 参数只能是 str 或 dict 类型")
        
        for key, value in kwargs.items():
            if key in self.config:
                self.config[key] = value
        

        # 如果使用地图，则忽略 figureRatio 参数，使用地图的宽高比
        if mapRange is not None and figureRatio is None:
            figureRatio = (mapRange[1] - mapRange[0]) / (mapRange[3] - mapRange[2])
        elif figureRatio is None:
            figureRatio = 1
            
        # 绘图区大小
        if figSize is None:
            figZoom = self.config['figZoom']
            fig_ratio = cols*figureRatio/rows
            figHeight = rows * figZoom / (self.config['sub_adj'][3] - self.config['sub_adj'][1])
            figWidth = rows * figZoom * fig_ratio / (self.config['sub_adj'][2] - self.config['sub_adj'][0])
            figSize = (figWidth, figHeight)
        
        self.figSize = figSize
        self.fig = plt.figure(figsize=figSize, dpi=self.config['dpi'])
        self.fig.subplots_adjust(*self.config['sub_adj'])
        self.rows = rows
        self.cols = cols

        if rows > 1 or cols > 1:
            self.use_gridspec(rows, cols, self.config['wspace'], self.config['hspace'])
            
            # 如果不是自定义子图，则生成规整的网格子图
            if not customAxes:
                self.axes = [[None for _ in range(cols)] for _ in range(rows)]
                for i in range(rows):
                    for j in range(cols):
                        ax = self.add_axes(i, j, mapRange=mapRange)
                        if not customMap and mapRange is not None:
                            self.draw_basemap(ax, onLeft=(j==0), onBottom=(i==rows-1))
        else:
            self.gridSpec = None
            self.add_axes(mapRange=mapRange)
            if not customMap and mapRange is not None:
                self.draw_basemap(self.axes[0][0], 
                                  onLeft=True, onBottom=True)

        self._contourf = None
        self._quiver = None

    def _load_config(self, path:str) -> None:
        """
        读取配置文件
        """
        with open(path, 'r') as f:
            self.config = yaml.load(f, Loader=yaml.FullLoader)
    
    def _cover_config(self, config: dict, kwargs: dict) -> dict:
        """
        覆盖配置文件中的参数，如果config中有则覆盖，没有则添加。
        """
        config = config.copy()
        for key, value in kwargs.items():
            config[key] = value
        return config

    def use_gridspec(self, rows: int, cols: int, 
                     wspace: float=0.1, hspace: float=0.1, **kwargs) -> None:
        """
        使用 GridSpec 布局子图,此函数会设置 self.gridSpec 。

        :param rows: 子图的行数。
        :param cols: 子图的列数。
        :param wspace: 子图的宽间隔。
        :param hspace: 子图的高间隔。
        """
        self.gridSpec = GridSpec(rows, cols, figure=self.fig, 
                                 wspace=wspace, hspace=hspace, **kwargs)
        
    def add_axes(self, row: Union[int, slice]=0, col: Union[int, slice]=0,
                 mapRange: list=None, proj: ccrs.Projection=None,
                 central_longitude: float=180, **kwargs) -> plt.Axes:
        """
        在已有的 GridSpec 框架中添加一个子图，并添加到 self.axes 列表中。
        如果是只有一张子图，没有使用 GridSpec，则直接添加一个全屏的子图(111)。

        例如在第2行插入一个宽3列的子图：`add_axes(1, slice(1, 4))`，等价`gridSpec[1, 1:4]`
        如果想在 slice 里去取到最后一个，可以使用 `slice(1, None)`

        :param row: 子图的行位置。可以用整数或slice来应用到 GridSpec 中。
        :param col: 子图的列位置。可以用整数或slice来应用到 GridSpec 中。
        :param mapRange: 地图的范围，如果设置了，则转为 cartopy 的 Axes 对象。
        :param proj: 地图的投影方式。不设置则使用默认的 PlateCarree 投影。
        :param central_longitude: 地图的中心经度，仅在 proj 不设置时有效。
        :param kwargs: 其他参数，传递给 matplotlib.figure.Figure.add_subplot()。
        """
        if proj is None:
            proj = ccrs.PlateCarree(central_longitude=central_longitude)

        if self.gridSpec is not None:
            if mapRange is not None:
                ax = self.fig.add_subplot(self.gridSpec[row, col], projection=proj, 
                                          **kwargs)
                ax.set_extent(mapRange, crs=ccrs.PlateCarree())
            else:
                ax = self.fig.add_subplot(self.gridSpec[row, col], **kwargs)
        else:
            if mapRange is not None:
                ax = self.fig.add_subplot(111, projection=proj, 
                                         **kwargs)
                ax.set_extent(mapRange, crs=ccrs.PlateCarree())
            else:
                ax = self.fig.add_subplot(111, **kwargs)
        
        self.axes[row][col] = ax
        return ax
        
    def draw_coastlines(self, ax: GeoAxes, **kwargs) -> None:
        """
        绘制海岸线，使用配置文件中的['coastlines']参数。
        传入的参数会覆盖配置文件中的参数。
        """
        params = self._cover_config(self.config['coastlines'], kwargs)
        ax.coastlines(**params)

    def draw_gridlines(self, ax: GeoAxes, crs=ccrs.PlateCarree(0), 
                       top=False, right=False, left=False, bottom=False,
                       **kwargs) -> Gridliner:
        """
        绘制网格线，使用配置文件中的['gridlines']参数。
        top, right, left, bottom 分别表示是否绘制上、右、左、下的标签，用于多图排列。
        """
        params = self._cover_config(self.config['gridlines'], kwargs)

        if 'xinterval' in params:
            xint = params['xinterval']
            del params['xinterval']
        else:
            xint = 10

        if 'yinterval' in params:
            yint = params['yinterval']
            del params['yinterval']
        else:
            yint = 10

        gl = ax.gridlines(crs=crs, 
                          xlocs=MultipleLocator(xint),
                          ylocs=MultipleLocator(yint),
                          **params)
        gl.top_labels = top
        gl.right_labels = right
        gl.left_labels = left
        gl.bottom_labels = bottom
        return gl

    def fill_land(self, ax: GeoAxes, **kwargs) -> None:
        """
        陆地填色，配置文件 fillLand 参数。
        """
        from cartopy.feature import LAND
        params = self._cover_config(self.config['fillLand'], kwargs)
        ax.add_feature(LAND, **params)

    def draw_basemap(self, ax: GeoAxes, onLeft=False, onBottom=False) -> None:
        """
        画基本地图地图，包括海岸线、网格线和陆地填色。根据配置文件的 'use*' 参数来决定是否绘制。
        onLeft, onBottom 分别表示是否绘制左、下的经纬度标签，用于多图排列。
        """
        if self.config['useCoastlines']:
            self.draw_coastlines(ax)
        if self.config['useGridlines']:
            if self.config['gridlinesLabelInEdge']:
                self.draw_gridlines(ax, left=onLeft, bottom=onBottom)
            else:
                self.draw_gridlines(ax, left=True, bottom=True)
        if self.config['useFillLand']:
            self.fill_land(ax)

    #TODO: draw_each, subplot_label, colorbar
    def for_each_axes(self, func:Callable[[plt.Axes, int, int, int, int], None],
                      useTqdm=True) -> None:
        """
        对每一个子图执行一个画图函数。

        :param func: 画图函数，接受四个参数。
            分别是当前子图的行、列、按行计数的序号（准确），按列计数的序号（只有是完整网格子图才准确）。
        :type func: Callable[[int, int, int, int], Union[QuadContourSet, None]]
        :param tqdm: 是否使用 tqdm 显示进度条。
        """
        k = 0
        if tqdm:
            pb = tqdm(total=self.rows*self.cols, desc="For each axes")
        for i in range(len(self.axes)):
            for j in range(len(self.axes[i])):
                c = func(self.axes[i][j], i, j, k, j*len(self.axes[i])+i)
                if c is not None:
                    self._contourf = c
                k += 1
                if useTqdm:
                    pb.update(1)
        return c

    def for_each_label(self, func:Callable[[int, int, int, int], str], **kwargs) -> None:
        """
        为每一个子图添加一个标签。

        :param func: 标签函数，接受四个参数。
            分别是当前子图的行、列、按行计数的序号（准确），按列计数的序号（只有是完整网格子图才准确）。
            返回一个字符串作为标签。
        :type func: Callable[[int, int, int, int], str]
        :param kwargs: 其他参数，传递给 matplotlib.pyplot.text()。默认使用配置 subLabel。
        """
        params = self._cover_config(self.config['subLabel'], kwargs)
        k = 0
        for i in range(len(self.axes)):
            for j in range(len(self.axes[i])):
                label = func(i, j, k, j*len(self.axes[i])+i)
                self.axes[i][j].text(s=label, 
                                     transform=self.axes[i][j].transAxes, **params)
                k += 1

    def colorbar(self, o :str='vertical',loc: List[float]=None, **kwargs) -> None:
        """
        添加一个 colorbar 在整个图的loc位置。会自动获取用 which 参数指定的子图的 contourf 对象。

        :param o: colorbar的方向，'vertical' 或 'horizontal'。
        :param loc: colorbar的位置，格式为 [left, bottom, width, height]。默认使用配置。
        :param which: 用于获取 contourf 对象的子图位置。

        """
        if loc is None:
            if o == 'vertical':
                loc = self.config['colorbarV']
            else:
                loc = self.config['colorbarH']
        
        if self._contourf is not None:
            if o == 'vertical':
                # 重新调整 fig的大小
                self.fig.set(figwidth=(2-loc[0])*self.figSize[0])
                self.fig.subplots_adjust(right=loc[0]-self.config['colorbarPad'])
                newAxes = self.fig.add_axes(loc)
                self.fig.colorbar(self._contourf, cax=newAxes, 
                                  orientation='vertical', **kwargs)
            else:
                # 重新调整 fig的大小
                self.fig.set(figheight=(1+loc[1]*2)*self.figSize[1])
                self.fig.subplots_adjust(bottom=loc[1]+loc[3])
                newAxes = self.fig.add_axes(loc)
                self.fig.colorbar(self._contourf, cax=newAxes, 
                                  orientation='horizontal', **kwargs)
        else:
            raise ValueError("没有 contourf 对象")
        
    
    def contourf(self, ax: Union[plt.Axes, GeoAxes], 
                 data: xr.DataArray, preset: str='none',
                 x: Union[np.ndarray, xr.DataArray]=None,
                 y: Union[np.ndarray, xr.DataArray]=None,
                 xName='longitude', yName='latitude', **kwargs) -> QuadContourSet:
        """
        画一个 contourf 图，使用配置文件中的 preset 参数。传入 DataArray 数据，
        使用 xName 和 yName 参数指定坐标。

        :param ax: 画图的子图对象。
        :param data: 画图的数据。
        :param x: 数据的x坐标，如果不传入则使用 data 的坐标。
        :param y: 数据的y坐标，如果不传入则使用 data 的坐标。
        :param xName: x轴的坐标名称。
        :param yName: y轴的坐标名称。
        :param preset: 配置文件中的预设，有效值：'t', 'q', 'z', 'vor', 'div', 'omega'。
        :param kwargs: 其他参数，传递给 matplotlib.pyplot.contourf()。

        :return: 返回 contourf 对象。
        """
        params = self._cover_config(self.config['contourf'][preset], kwargs)
        if isinstance(params['levels'], list):
            params['levels'] = np.linspace(*params['levels'])

        if x is None:
            x = data[xName]
        if y is None:
            y = data[yName]

        c = ax.contourf(x, y, data, 
                        transform=ccrs.PlateCarree() if isinstance(ax, GeoAxes) else None,
                        **params)
        self._contourf = c
        return c
    
    def quiver(self, ax: Union[plt.Axes, GeoAxes], 
               u: xr.DataArray, v: xr.DataArray, lowRes: int = 4,
               xName='longitude', yName='latitude', **kwargs) -> Quiver:
        """
        画一个 quiver 图。传入 u 和 v 的 DataArray 数据，使用 xName 和 yName 参数指定坐标。

        :param ax: 画图的子图对象。
        :param u: x方向的数据。
        :param v: y方向的数据。
        :param xName: x轴的坐标名称。
        :param yName: y轴的坐标名称。
        :param kwargs: 其他参数，传递给 matplotlib.pyplot.quiver()，默认使用配置。

        :return: 返回 quiver 对象。
        """
        params = self._cover_config(self.config['quiver'], kwargs)
        q = ax.quiver(u[xName].values[::lowRes], u[yName].values[::lowRes], 
                  u.values[::lowRes, ::lowRes], v.values[::lowRes, ::lowRes], 
                  transform=ccrs.PlateCarree() if isinstance(ax, GeoAxes) else None, 
                  **params)
        self._quiver = q
        return q
        
    def quiverkey(self, ax: Union[plt.Axes, GeoAxes]=None, quiver: Quiver=None,
                  x: float=0.9, y: float=0.97, length:float = 5, **kwargs) -> None:
        """
        添加一个 quiver 的 key，使用配置文件中的 quiverkey 参数。

        :param ax: 画图的子图对象。
        :param kwargs: 其他参数，传递给 matplotlib.pyplot.quiverkey()，默认使用配置。
        """
        params = self._cover_config(self.config['quiverkey'], kwargs)

        if quiver is None:
            quiver = self._quiver
        
        if 'unit' in params:
            unit = params['unit']
            del params['unit']
        else:
            unit = ""
        
        if ax is None:
            ax = self.fig.gca()
        ax.quiverkey(quiver, x, y, length, "%s %s" % (str(length), unit), **params)

    def fig_title(self, title: str, fontsize=12, **kwargs) -> None:
        """
        添加总标题。

        :param title: 标题。
        :param fontsize: 字体大小。
        :param kwargs: 其他参数，传递给 matplotlib.pyplot.title()。
        """
        self.fig.suptitle(title, fontsize=fontsize, **kwargs)